package registry

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/Unixeno/TheMoon/logging"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

var logger = logging.GetLogger("registry")

// Config registry config struct
type Config struct {
	Endpoint       string
	RegistryDomain string
	Proxy          string
	Insecure       bool
	Username       string
	Password       string
	RepoName       string
	RepoTag        string
	SavePath       string
}

// Registry registry object
type Registry struct {
	URL          *url.URL
	Host         string
	RegistryHost string
	client       *http.Client
	Auth         auth
	Config       Config
}

type auth struct {
	Token       string    `json:"token,omitempty"`
	AccessToken string    `json:"access_token,omitempty"`
	ExpiresIn   int       `json:"expires_in,omitempty"`
	IssuedAt    time.Time `json:"issued_at,omitempty"`
}

// Manifests is the image manifest struct
type Manifest struct {
	Config        ManifestConfig  `json:"config,omitempty"`
	Layers        []ManifestLayer `json:"layers,omitempty"`
	MediaType     string          `json:"mediaType,omitempty"`
	SchemaVersion int             `json:"schemaVersion"`
}
type ManifestConfig struct {
	Digest    string `json:"digest"`
	MediaType string `json:"mediaType"`
	Size      int    `json:"size"`
}

type ManifestLayer struct {
	Digest    string `json:"digest"`
	MediaType string `json:"mediaType"`
	Size      int    `json:"size"`
}

type ManifestList struct {
	Manifests     []ManifestInfo `json:"manifests,omitempty"`
	MediaType     string         `json:"mediaType,omitempty"`
	SchemaVersion int            `json:"schemaVersion"`
}

type ManifestInfo struct {
	Digest    string       `json:"digest"`
	MediaType string       `json:"mediaType"`
	Size      int          `json:"size"`
	Platform  PlatformInfo `json:"platform"`
}

type PlatformInfo struct {
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
}

func getProxy(proxy string) func(*http.Request) (*url.URL, error) {
	if len(proxy) > 0 {
		proxyURL, _ := url.Parse(proxy)
		return http.ProxyURL(proxyURL)
	}
	return http.ProxyFromEnvironment
}

// TokenExpired returns whether or not an auth token has expired
func (reg *Registry) TokenExpired() bool {
	duration := time.Since(reg.Auth.IssuedAt)
	if int(duration.Seconds()) > reg.Auth.ExpiresIn {
		logger.Info("auth token expired")
		return true
	}
	return false
}

// New creates a new Registry object
func New(rc Config) (*Registry, error) {
	u, e := url.Parse(rc.Endpoint)
	if e != nil {
		return nil, e
	}
	if u.Host == "" {
		u.Host = u.Path
	}
	origURL := u
	host := fmt.Sprintf("%s://%s", u.Scheme, u.Host)
	registryHost := ""
	if rc.RegistryDomain != "" {
		u, e = url.Parse(rc.RegistryDomain)
		if e != nil {
			return nil, e
		}
		if u.Host == "" {
			u.Host = u.Path
		}
		registryHost = fmt.Sprintf("%s://%s", u.Scheme, u.Host)
	}
	client := &http.Client{
		Transport: &http.Transport{
			Proxy:           getProxy(rc.Proxy),
			TLSClientConfig: &tls.Config{InsecureSkipVerify: rc.Insecure},
		},
	}
	return &Registry{
		URL:          origURL,
		Host:         host,
		RegistryHost: registryHost,
		client:       client,
		Config:       rc}, nil
}

// GetToken get a docker registry API pull token
func (reg *Registry) GetToken() error {
	u := fmt.Sprintf("https://auth.docker.io/token?service=registry.docker.io&scope=repository:%s:pull", reg.Config.RepoName)
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	if reg.Config.Username != "" && reg.Config.Password != "" {
		req.SetBasicAuth(reg.Config.Username, reg.Config.Password)
	}
	res, err := reg.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return fmt.Errorf("HTTP Error: %s", res.Status)
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}

	var a = new(auth)

	err = json.Unmarshal(body, &a)
	if err != nil {
		return err
	}
	reg.Auth = *a

	return nil
}

func (reg *Registry) doGet(url string, headers map[string]string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", reg.Auth.Token))
	// add additional headers
	if headers != nil {
		for key, value := range headers {
			req.Header.Add(key, value)
		}
	}
	res, err := reg.client.Do(req)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		res.Body.Close()
		return nil, fmt.Errorf("HTTP Error: %s", res.Status)
	}
	return res, nil
}

func (reg *Registry) GetManifestList() (string, error) {
	headers := make(map[string]string)
	headers["Accept"] = "application/vnd.docker.distribution.manifest.list.v2+json"
	if reg.TokenExpired() {
		reg.GetToken()
	}
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", reg.Host, reg.Config.RepoName, reg.Config.RepoTag)
	res, err := reg.doGet(url, headers)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()
	rawJSON, err := ioutil.ReadAll(res.Body)
	manifestList := ManifestList{}
	if err := json.Unmarshal(rawJSON, &manifestList); err != nil {
		return "", err
	}

	for _, m := range manifestList.Manifests {
		if m.Platform.Architecture == runtime.GOARCH && m.Platform.OS == "linux" {
			return m.Digest, nil
		}
	}

	return "", fmt.Errorf("invalid Architecture")
}

// ReposManifest gets docker image manifest for name:tag
func (reg *Registry) ReposManifest(manifestDigest string) (*Manifest, error) {
	headers := make(map[string]string)
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", reg.Host, reg.Config.RepoName, manifestDigest)
	fmt.Println("manifest url ", url)
	headers["Accept"] = "application/vnd.docker.distribution.manifest.v2+json"
	if reg.TokenExpired() {
		reg.GetToken()
	}

	res, err := reg.doGet(url, headers)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	rawJSON, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	m := new(Manifest)
	if err := json.Unmarshal(rawJSON, &m); err != nil {
		return nil, err
	}

	return m, nil
}

// RepoGetConfig gets docker image config JSON
func (reg *Registry) RepoGetConfig(manifest *Manifest) (string, error) {
	// Create the file
	configFilePath := filepath.Join(reg.Config.SavePath, fmt.Sprintf("%s.json", strings.TrimPrefix(manifest.Config.Digest, "sha256:")))
	out, err := os.Create(configFilePath)
	if err != nil {
		return "", fmt.Errorf("create config file failed")
	}
	defer out.Close()
	// Download config
	headers := make(map[string]string)
	url := fmt.Sprintf("%s/v2/%s/blobs/%s", reg.Host, reg.Config.RepoName, manifest.Config.Digest)
	headers["Accept"] = manifest.Config.MediaType

	if reg.TokenExpired() {
		reg.GetToken()
	}

	logger.WithField("filepath", configFilePath).Debug("downloading...")
	res, err := reg.doGet(url, headers)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	// Write the body to file
	_, err = io.Copy(out, res.Body)
	if err != nil {
		return "", fmt.Errorf("writing config file failed")
	}

	return filepath.Base(configFilePath), nil
}

func getTarFileSuffix(mediaType string) string {
	t := strings.Split(mediaType, ".")
	s := t[len(t)-1]

	switch s {
	case "gzip":
		return "tar.gz"
	case "zstd":
		return "tar.zst"
	default:
		return "tar"
	}
}

// RepoGetLayers gets docker image layer tarballs
func (reg *Registry) RepoGetLayers(layers []ManifestLayer) ([]string, error) {
	var layerFiles []string

	for _, layer := range layers {

		// Create the TAR file
		suffix := getTarFileSuffix(layer.MediaType)
		tarFilePath := filepath.Join(reg.Config.SavePath, fmt.Sprintf("%s.%s", strings.TrimPrefix(layer.Digest, "sha256:"), suffix))
		layerFiles = append(layerFiles, filepath.Base(tarFilePath))
		out, err := os.Create(tarFilePath)
		if err != nil {
			return nil, fmt.Errorf("create store file failed")
		}
		defer out.Close()

		// Download layer
		headers := make(map[string]string)
		url := fmt.Sprintf("%s/v2/%s/blobs/%s", reg.Host, reg.Config.RepoName, layer.Digest)
		headers["Accept"] = layer.MediaType

		if reg.TokenExpired() {
			reg.GetToken()
		}
		logger.WithField("filepath", tarFilePath).WithField("size", layer.Size).Debug("downloading...")
		res, err := reg.doGet(url, headers)
		if err != nil {
			return nil, err
		}
		defer res.Body.Close()
		// Write the body to file
		_, err = io.Copy(out, res.Body)
		if err != nil {
			return nil, fmt.Errorf("write tar file failed")
		}
	}
	logger.Info("finish download")
	return layerFiles, nil
}
